# 参数：数据宽度data_width、数据深度data_depth、en默认1？
# 先只用ramb18e1，也就是width最大16，depth 1k？
# INIT_00到INIT_3F，每个256位,16进制

# yosys好像不支持reg的input output，只支持wire类型

from base.port_item import PortItem
from base.design import Design
from base.module import Module
from base.net import Net

# from cells.ramb18e1 import RAMB18E1
from cells.ramb36e1 import RAMB36E1

import json
from base import json_util

from base import util
class RomIP:
    

    def __init__(self,data_width:int=4, data_depth:int=2, module_name="rom_ip", coe_path=None) -> None:
        self.ram_mode="TDP"
        self.port_clk:PortItem = None
        self.port_en:PortItem = None
        self.port_addr:PortItem = None
        self.port_dout:PortItem = None
        self.design:Design = None
        self.top_module: Module = None

        # 存储 coe 文件数据
        self.init_radix = 16
        self.init_data = []
        
        # index 对应 ramb 的 index ，value 是该 ramb 的 init数据，结构如下
        '''
            [
                {
                    "cur_bit_count": 14,
                    "cur_loc": 15,
                    "INIT_00" : "",
                    "INIT_01" : "",
                    ...
                },
                {
                    "cur_bit_count": 14,
                    "cur_loc": 2,
                    "INIT_00" : "",
                    "INIT_01" : "",
                    ...
                },
                ...
            ]

        '''
        self.init_data_to_ramb = []

        # 暂时只支持18位宽、18k容量
        # assert(data_width in [1,2,4,8,16,32])
        # assert(data_width*data_depth<=32*1024)

        self.data_width = data_width
        self.data_depth = data_depth
        self.module_name = module_name

        # 计算需要多少个ramb36e1 实现数据宽度,并联相关名称带row
        self.count_row = math.ceil(data_width/32)
        # 计算单个ramb36e1的数据宽度
        self.single_data_width = data_width
        if data_width>32:
            self.single_data_width = 32

        # 计算出ramb36e1的地址线数量，暂时不考虑串联
        self.addr_width = get_addr_width(self.data_depth)
        self.addr_move_count = get_addr_move_count(self.single_data_width)

        # # 根据参数计算输入输出参数
        self.ram_data_width = get_read_width_parameter(self.single_data_width) #用来传递parameter

        if coe_path is not None:
            self.set_init_from_coe_path(coe_path)
        print(self.init_data)
        
        self.distribute_init_data()
        pass

    # def get_inits(self):
    #     total_mem = self.data_depth*self.data_width
    # def set_init_value(self,init_values:str):
    #     # 每个INIT_XX为256位，即64个16进制数字
    #     # 16进制,低位在右边
        
    #     pass
    
    # 根据设置的参数，生成module内容
    def gen_module(self):
        self.design = Design()
        self.top_module = Module(self.module_name)

        # 第1步，添加netnames
        # clk
        net_clk = Net("clk",1)
        self.top_module.add_net(net_clk)
        # addr
        net_addr = Net("addr",self.addr_width)
        self.top_module.add_net(net_addr)
        # dout
        net_dout = Net("dout", self.data_width)# ramb36e1的所有DO没有用的，必须要接到空net上
       
        self.top_module.add_net(net_dout)


        # 第二步，添加ports，并将port的bits设置成相关netnames的bits，表达连接关系
        port_clk = PortItem(name="clk")
        port_clk.get_pin().set_bit(net_clk.get_bit())
        self.top_module.add_port(port_clk)
        port_addr = PortItem(width=self.addr_width,name="addr")
        for index,pin in enumerate(port_addr.get_all_pins()):
            pin.set_bit(net_addr.get_bit(index))
        self.top_module.add_port(port_addr)
        port_dout = PortItem(type="output",width=self.data_width,name="dout")
        for index,pin in enumerate(port_dout.get_all_pins()):
            pin.set_bit(net_dout.get_bit(index))
        self.top_module.add_port(port_dout)

        # 第三步，添加cells，并指定connections，表达连接关系
        dout_index = 0
        dout_full = False
        for index in range(self.count_row):
            cell_ramb36e1_item = RAMB36E1("cell_ramb36e1_"+str(index))
            cell_ramb36e1_item.add_connection(port_clk.get_pin(),"CLKARDCLK")
            for j in range(self.addr_width):
                cell_ramb36e1_item.add_connection(port_addr.get_pin_by_index(j),"ADDRARDADDR",j+self.addr_move_count)
            # dout连接,根据算出来的单个ramb36e1数据宽度添加，
            for j in range(self.single_data_width):
                cell_ramb36e1_item.add_connection(port_dout.get_pin_by_index(dout_index),"DOADO",j)
                dout_index+=1
                if dout_index == self.data_width:
                    dout_full = True
                    break
            if dout_full: # dout所有引脚分配完成，ramb36多余引脚接到空net上
                self.top_module.set_cell_port_empty(cell_ramb36e1_item,"DOADO",list(range(j+1,32)))
            # wea要连接低电平
            for j in range(4):
                cell_ramb36e1_item.add_connection("0","WEA",j)
                
            # ena接高电平
            cell_ramb36e1_item.add_connection("1","ENARDEN")

            cell_ramb36e1_item.set_read_width_a(self.ram_data_width)
            # 设置工作模式
            cell_ramb36e1_item.set_ram_mode("TDP")
            
            # 根据 index 获取每个 ramb 的 init 数据，并通过循环赋值
            cur_ramb_init_group = self.init_data_to_ramb[index]
            cur_ramb_init_group_len = cur_ramb_init_group['cur_loc']+1
            for init_index in range(cur_ramb_init_group_len):
                init_name = f"INIT_{init_index:02X}"
                if init_name in cur_ramb_init_group and len(cur_ramb_init_group[init_name]) > 0 :
                    cell_ramb36e1_item.set_parameter(f"INIT_{init_index:02X}",cur_ramb_init_group[init_name].rjust(256, '0'))
            self.top_module.add_cell(cell_ramb36e1_item)
        
        # 循环结束

        pass

    # 将 INIT_data 依照 ramb 的宽度进行分配
    def distribute_init_data(self):
        
        for i in range(self.count_row):
            # 创建 i 个 字典，存储每个 ramb 的 init
            self.init_data_to_ramb.append({
                "cur_bit_count": 0,
                "cur_loc": 0,
                "INIT_00" : ""
            })
        
        for init_data in self.init_data:
            end_idx = len(init_data)
            for cur_ramb in self.init_data_to_ramb:
                # 使用校验位时，每次需要处理的宽度就是read_width(data+parity)
                # 不使用校验位时，每次需要处理的宽度就是data_width
                per_width = self.single_data_width
                start_idx = max(end_idx - per_width, 0)
                if start_idx == 0:
                    # 如果 init_data 剩余长度不足以填满 per_width，则填充 '0'
                    per_data = init_data[start_idx:end_idx].rjust(per_width, '0')
                else:
                    per_data = init_data[start_idx:end_idx]
                
                # 获取当前的 init_xx
                init_key = f"INIT_{cur_ramb['cur_loc']:02X}"
                # 当前 init_xx 剩余的可用位数
                remaining_bits = 256 - cur_ramb['cur_bit_count']
                cur_ramb[init_key] = per_data + (cur_ramb[init_key] if init_key in cur_ramb else "")
                cur_ramb["cur_bit_count"] += len(per_data)
                
                # 如果当前 init_xx 达到 256 位，重置 cur_bit_count 并切换到下一个 init_xx
                if cur_ramb["cur_bit_count"] >= 256:
                    cur_ramb["cur_bit_count"] = 0
                    cur_ramb["cur_loc"] += 1

                end_idx = start_idx
                
                if end_idx == 0:
                    break  # 当前 init_data 已全部分配完，处理下一个 init_data
        
    # 根据传入的coe文件，得到INIT值
    def set_init_from_coe_path(self,coe_path:str):
        
        self.init_radix, init_data = util.read_coe_file(coe_path)
        
        # 目标宽度和深度
        target_depth = self.data_depth
        target_width = self.data_width
        
        # 初始化结果列表
        processed_data = []
        init_values = init_data.split(',')
        
        # 处理 init_data 的长度不足或超出的情况
        init_len = len(init_values) # init数据传入depth
        if init_len < target_depth:
            # 不够则补零
            # init_values.extend(['0'] * (target_depth - init_len))
            # 这里暂时不填充0 ，填充0后会导致init太大，对实际结果又没有正向的反馈
            pass
        elif init_len > target_depth:
            # 超出则截断后面的数据
            init_values = init_values[:target_depth]
        
        rang_depth = target_depth if init_len > target_depth else init_len
        
        for depth_index in range(rang_depth):
            # 取出当前深度的数据
            row_data = init_values[depth_index]
            # 将十六进制字符串转换为二进制字符串
            bin_data = util.hex_to_bin(row_data)
            # 调整到指定宽度
            if len(bin_data) > target_width:
                # 宽度大于 width 时，截取低位部分
                bin_data = bin_data[-target_width:]
            else:
                # 宽度小于等于 width 时，用 0 填充高位部分
                bin_data = bin_data.zfill(target_width)
        
            # 根据输出格式添加到结果列表
            processed_data.append(bin_data)

        self.init_data = processed_data
        

    # 将网表输出为json文件
    def to_json(self):
        out_data = dict()
        out_data["creator"] = "HybrdChip"
        m = dict()
        m["attributes"] = {"hdlname":"rom_ip"}

        # netnames
        netnames = {}
        for net in self.top_module.get_all_nets():
            netnames[net.get_name()] = net.to_dict()
        m["netnames"] = netnames
        


        # ports
        m["ports"] = dict()
        for port in self.top_module.get_all_ports():
            m["ports"][port.get_name()] = {"direction":port.get_type(),"bits":[]}
            for pin in port.get_all_pins():
                m["ports"][port.get_name()]["bits"].append(pin.get_bit())

        # cells
        m["cells"] = dict()
        for cname,cell in self.top_module.get_all_cells().items():
            m["cells"][cname] = cell
            

        out_data["modules"] = dict()
        out_data["modules"][self.top_module.get_name()] = m
        ss = json.dumps(out_data,default=json_util.json_encoder,indent=4)
        return ss



        
import math

# 根据数据深度，计算module的地址线位数
def get_addr_width(depth:int):
    if depth < 1:
        return 0
    elif depth > 1024:
        return 
    return math.ceil(math.log2(depth))

# 根据计算的ramb36数据宽度，计算低几位addr不用
def get_addr_move_count(single_data_width):
    if single_data_width==1:
        return 0
    elif single_data_width == 2:
        return 1
    elif single_data_width ==4:
        return 2
    elif single_data_width == 8:
        return 3
    elif single_data_width ==16:
        return 4
    elif single_data_width == 32:
        return 5

# 根据计算的ramb36数据宽度，得到parameter里设置的read_width_a
def get_read_width_parameter(x):
    if x == 1:
        return 1
    elif x == 2:
        return 2
    elif x <= 4:
        return 4
    elif x <= 8:
        return 9
    elif x <= 16:
        return 18
    else:
        return 36
    


if __name__ == "__main__":
    rom_ip_inst = RomIP(data_depth=4,data_width=64,coe_path="/home/zhaoshuai/Source/ip_generator_py/src/data/rom.coe")
    rom_ip_inst.gen_module()
    with open("out.json",'w') as f:
        f.write(rom_ip_inst.to_json())
    pass
    # rom_ip_inst.port_clk=PortItem(name="clk")
    # rom_ip_inst.port_addr=PortItem(name="addr",width=4)
    # rom_ip_inst.port_en=PortItem(name="en")
    # rom_ip_inst.port_dout=PortItem(name="dout",width=4)

    # s = rom_ip_inst.write_verilog()
    # with open("out.v",'w') as f:
    #     f.write(s)